#include "encrypto.h"
#include <assert.h>
#include <chrono>
#include <iostream>
#include <openssl/aes.h>
#include <openssl/evp.h>
#include <openssl/objects.h>
#include <stdio.h>
#include <string>

using namespace std;

CEncryptyFile::CEncryptyFile() { }

CEncryptyFile::~CEncryptyFile() { }

// AES 加密算法
// 注意：参数和返回值全部是二进制数据
std::string CEncryptyFile::AesEncrypt(const std::string& source, const std::string& key)
{
    // EVP_des_ecb  EVP_aes_128_ecb,AES/DES两种算法
    EVP_CIPHER_CTX* ctx = EVP_CIPHER_CTX_new();
    EVP_CIPHER_CTX_init(ctx);
    int ret = EVP_EncryptInit_ex(ctx, EVP_aes_128_ctr(), NULL, (const unsigned char*)key.c_str(), NULL);
    if (ret != 1) {
        cout << "error while EVP_EncryptInit_ex" << endl;
        return "";
    }

    // unsigned char* result = new unsigned char[source.length() + 64]; // 弄个足够大的空间
    string result;
    result.resize(source.length() + AES_BLOCK_SIZE);
    int len1 = 0;

    auto pos1 = chrono::steady_clock::now();

    ret = EVP_EncryptUpdate(ctx, (unsigned char*)result.c_str(), &len1, (const unsigned char*)source.c_str(), source.length());
    if (ret != 1) {
        cout << "error while EVP_EncryptUpdate: " << ret << ", len1: " << len1 << endl;
        // delete[] result;
        return "";
    }

    int len2 = 0;
    ret = EVP_EncryptFinal_ex(ctx, (unsigned char*)result.c_str() + len1, &len2);
    if (1 != ret) {
        EVP_CIPHER_CTX_cleanup(ctx);
        // delete[] result;
        return "";
    }

    auto pos2 = chrono::steady_clock::now();
    chrono::duration<double, milli> dur1 = pos2 - pos1;
    cout << "encrypt [" << source.length() << "]B's data cost: " << dur1.count() << " ms" << endl;

    ret = EVP_CIPHER_CTX_cleanup(ctx);
    if (1 != ret) {
        EVP_CIPHER_CTX_cleanup(ctx);
        // delete[] result;
        return "";
    }

    std::string res((char*)result.c_str(), len1 + len2);
    // delete[] result;
    return res;
}

// 解密算法
std::string CEncryptyFile::AesDecrypt(const std::string& ciphertext, const std::string& key)
{
    // EVP_des_ecb  EVP_aes_128_ecb,AES/DES两种算法
    EVP_CIPHER_CTX* ctx = EVP_CIPHER_CTX_new();
    EVP_CIPHER_CTX_init(ctx);
    int ret = EVP_DecryptInit_ex(ctx, EVP_aes_128_ctr(), NULL, (const unsigned char*)key.data(), NULL);
    if (1 != ret)
        return "";

    // unsigned char* result = new unsigned char[ciphertext.length() + 64]; // 弄个足够大的空间
    string result;
    result.resize(ciphertext.length() + AES_BLOCK_SIZE);

    auto pos1 = chrono::steady_clock::now();

    int len1 = 0;
    ret = EVP_DecryptUpdate(ctx, (unsigned char*)result.c_str(), &len1, (const unsigned char*)ciphertext.data(), ciphertext.length());
    if (1 != ret) {
        EVP_CIPHER_CTX_cleanup(ctx);
        // delete[] result;
        return "";
    }

    int len2 = 0;
    ret = EVP_DecryptFinal_ex(ctx, (unsigned char*)result.c_str() + len1, &len2);
    if (1 != ret) {
        EVP_CIPHER_CTX_cleanup(ctx);
        // delete[] result;
        return "";
    }

    auto pos2 = chrono::steady_clock::now();
    chrono::duration<double, milli> dur1 = pos2 - pos1;
    cout << "decrypt [" << ciphertext.length() << "]B's data cost: " << dur1.count() << " ms" << endl;

    ret = EVP_CIPHER_CTX_cleanup(ctx);
    if (1 != ret) {
        // delete[] result;
        return "";
    }
    std::string res((char*)(unsigned char*)result.c_str(), len1 + len2);
    // cout << "decrypt len1: " << len1 << ", len2: " << len2 << endl;
    // delete[] result;
    return res;
}

std::string CEncryptyFile::AesEncryptorDecryptFile(std::string strFilePath, std::string key, std::string strOutFilePath, bool bEncrypt)
{
    FILE* fp = NULL;
    if ((fp = fopen(strFilePath.c_str(), "rb")) == nullptr) {
        return "";
    }

    fseek(fp, 0L, SEEK_END);
    const int length = ftell(fp);
    fseek(fp, 0, SEEK_SET);

    char* buf = new char[length + 64];

    int bytes = fread(buf, 1, length, fp);
    fclose(fp);
    fp = NULL;

    std::string source(buf, bytes); // 二进制数据

    // 加密
    //	std::string keyData(key, key.length());
    std::string keyData = key;

    std::string enc;
    if (bEncrypt) {
        enc = AesEncrypt(source, keyData);
    } else {
        enc = AesDecrypt(source, keyData);
    }

    if (enc.empty()) {
        return "";
    }

    // 输出到文件
    if (!strOutFilePath.empty()) {
        if ((fp = fopen(strOutFilePath.c_str(), "wb")) == nullptr) {
            fwrite(enc.data(), 1, enc.length(), fp);
            fclose(fp);
            fp = NULL;
        }
    }

    return enc;
}

int dataLen = 1024 * 1024 * 4;
const int roundNum = 1;
const std::string key = "0123456789123456";

int main(int argc, char* argv[])
{
    if (argc != 2) {
        std::cout << "err num of args" << std::endl;
        return -1;
    }
    dataLen = atoi(argv[1]);

    CEncryptyFile ce;
    std::string in;
    in.resize(dataLen);
    fill(in.begin(), in.end(), 'a');
    size_t loop_num = 21;
    for (size_t i = 0; i < loop_num; i++) {
        std::string out = ce.AesEncrypt(in, key);
        if (out.length() == 0)
            cout << 0 << endl;
        // for (int i = 0; i < dataLen; i++) {
        //     if (in[i] != out[i]) {
        //         cout << "out data is not equal to in data" << endl;
        //         break;
        //     }
        // }
        std::string inRecoverd = ce.AesDecrypt(out, key);
        // cout << "length of out: " << out.length() << ", data recoverd: " << inRecoverd.length() << endl;
        // if (inRecoverd.length() < dataLen)
        //     cout << "the length of data recoverd: " << inRecoverd.length() << endl;
        // for (int i = 0; i < dataLen; i++) {
        //     if (in[i] != inRecoverd[i]) {
        //         cout << "data is not equal" << endl;
        //         break;
        //     }
        // }
    }

    return 0;
}